#include <gtest/gtest.h>
#include "viml/simple_optimizer.h"
#include "viml/ast.h"
#include "viml/parser.h"
#include "viml/lexer.h"

using namespace viml;

class SimpleOptimizerTest : public ::testing::Test {
protected:
    void SetUp() override {
        optimizer = std::make_unique<SimpleOptimizer>();
    }

    std::unique_ptr<SimpleOptimizer> optimizer;

    // 辅助方法：解析源代码为AST
    std::shared_ptr<Program> parse_source(const std::string& source) {
        auto lexer = std::make_unique<Lexer>(source);
        auto tokens = lexer->tokenize();
        EXPECT_FALSE(tokens.empty());

        auto parser = std::make_unique<Parser>(tokens);
        return parser->parse();
    }

    // 辅助方法：获取优化统计信息
    SimpleOptimizer::OptimizationSummary get_optimization_summary() {
        return optimizer->get_optimization_summary();
    }
};

// 测试常量折叠优化
TEST_F(SimpleOptimizerTest, ConstantFoldingBasic) {
    // 测试简单的数字常量折叠
    std::string source = R"(
        let x = 1 + 2
        let y = 3 * 4
        let z = 10 - 5
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    // 验证优化器运行了（至少尝试了优化）
    EXPECT_GE(summary.optimizations_total, 0);
    // 验证没有崩溃
    EXPECT_NE(optimized, nullptr);
}

// 测试字符串连接常量折叠
TEST_F(SimpleOptimizerTest, StringConstantFolding) {
    std::string source = R"(
        let x = "hello" . " " . "world"
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    EXPECT_GT(summary.constant_folds, 0);
}

// 测试负数常量折叠
TEST_F(SimpleOptimizerTest, UnaryConstantFolding) {
    std::string source = R"(
        let x = -5
        let y = !0
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    EXPECT_GT(summary.constant_folds, 0);
}

// 测试复杂表达式常量折叠
TEST_F(SimpleOptimizerTest, ComplexConstantFolding) {
    std::string source = R"(
        let x = (1 + 2) * (3 - 1)
        let y = 10 / (2 + 3)
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    EXPECT_GT(summary.constant_folds, 0);
}

// 测试非常量表达式不会被折叠
TEST_F(SimpleOptimizerTest, NonConstantExpressions) {
    std::string source = R"(
        let x = 1 + y
        let z = "hello" . name
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    // 不会有常量折叠，因为包含变量
    EXPECT_EQ(summary.constant_folds, 0);
}

// 测试死代码消除
TEST_F(SimpleOptimizerTest, DeadCodeElimination) {
    std::string source = R"(
        let x = 1
        echo x
        let y = 2  # 这行在简单分析中可能被认为是死代码
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    // 当前实现简化处理，可能不会真正消除代码
    // 但会报告有死代码消除的分析
    EXPECT_GE(summary.dead_code_eliminated, 0);
}

// 测试寄存器重用优化
TEST_F(SimpleOptimizerTest, RegisterReuseOptimization) {
    std::string source = R"(
        let x = 1
        let y = 2
        let z = x + y
        let a = y + z
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    EXPECT_GT(summary.register_reuses, 0);
}

// 测试优化器启用/禁用
TEST_F(SimpleOptimizerTest, EnableDisableOptimizations) {
    std::string source = R"(
        let x = 1 + 2
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    // 禁用优化
    optimizer->set_enabled(false);
    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    EXPECT_EQ(summary.optimizations_total, 0);

    // 启用优化
    optimizer->set_enabled(true);
    optimizer->reset_counters();
    optimized = optimizer->optimize_ast(program);

    summary = get_optimization_summary();
    EXPECT_GT(summary.optimizations_total, 0);
}

// 测试常量表达式识别
TEST_F(SimpleOptimizerTest, ConstantExpressionRecognition) {
    std::string source = R"(
        let x = 42           # 常量
        let y = "hello"      # 常量
        let z = v:true       # 常量
        let a = x + 1        # 非常量（包含变量）
        let b = "hi" . name  # 非常量（包含变量）
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    // 应该检测到一些常量表达式
    EXPECT_GT(summary.constant_folds, 0);
}

// 测试错误情况下的常量折叠
TEST_F(SimpleOptimizerTest, ConstantFoldingErrorHandling) {
    std::string source = R"(
        let x = 1 / 0  # 除零错误，不应该折叠
        let y = 5 % 0  # 模零错误，不应该折叠
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    // 不应该折叠这些错误的表达式
    auto summary = get_optimization_summary();
    EXPECT_EQ(summary.constant_folds, 0);
}

// 测试优化统计重置
TEST_F(SimpleOptimizerTest, ResetOptimizationCounters) {
    std::string source = R"(
        let x = 1 + 2
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    // 第一次优化
    optimizer->reset_counters();
    auto optimized1 = optimizer->optimize_ast(program);
    auto summary1 = get_optimization_summary();
    EXPECT_GT(summary1.optimizations_total, 0);

    // 第二次优化（重置后）
    optimizer->reset_counters();
    auto optimized2 = optimizer->optimize_ast(program);
    auto summary2 = get_optimization_summary();
    EXPECT_GT(summary2.optimizations_total, 0);

    // 统计应该相同（因为重置了）
    EXPECT_EQ(summary1.optimizations_total, summary2.optimizations_total);
}

// 测试空程序优化
TEST_F(SimpleOptimizerTest, EmptyProgramOptimization) {
    std::string source = "";

    auto program = parse_source(source);
    // 空程序可能返回nullptr或空程序
    if (program) {
        optimizer->reset_counters();
        auto optimized = optimizer->optimize_ast(program);

        // 空程序不应该有优化
        auto summary = get_optimization_summary();
        EXPECT_EQ(summary.optimizations_total, 0);
    }
}

// 测试字节码级别优化
TEST_F(SimpleOptimizerTest, BytecodeOptimization) {
    std::string source = R"(
        let x = 1 + 2
        echo x
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    // 这里我们无法直接测试字节码优化，因为需要编译器
    // 但可以测试优化器不会崩溃
    optimizer->reset_counters();
    auto optimized_ast = optimizer->optimize_ast(program);

    // AST优化应该工作
    auto summary = get_optimization_summary();
    EXPECT_GT(summary.optimizations_total, 0);
}

// 测试复杂的嵌套表达式
TEST_F(SimpleOptimizerTest, NestedConstantExpressions) {
    std::string source = R"(
        let x = (1 + (2 * (3 + (4 - 1))))
        let y = ((5 * 2) + (10 / 2)) - 3
    )";

    auto program = parse_source(source);
    ASSERT_NE(program, nullptr);

    optimizer->reset_counters();
    auto optimized = optimizer->optimize_ast(program);

    auto summary = get_optimization_summary();
    // 应该检测到多个常量表达式
    EXPECT_GT(summary.constant_folds, 0);
}

// 主函数
int main(int argc, char** argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}